package com.cn.lp.ai.factory.domain.vector;

import com.cn.lp.ai.collection.AiCollection;
import com.cn.lp.ai.common.BusinessException;
import com.cn.lp.ai.embedding.AiEmbeddingService;
import com.cn.lp.ai.embedding.AiEmbeddingServiceHandler;
import com.cn.lp.ai.model.AiModel;
import com.cn.lp.ai.tag.AiTag;
import com.cn.lp.ai.tag.BaseAiModelTag;
import com.cn.lp.ai.vector.AiVectorService;
import com.cn.lp.ai.vector.VectorType;
import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.RedisVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.boot.autoconfigure.data.redis.RedisConnectionDetails;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import redis.clients.jedis.JedisPooled;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

@Service
@AllArgsConstructor
@Configuration
public class RedisVectorService implements AiVectorService {

    private final RedisConnectionDetails redisConnectionDetails;
    private AiEmbeddingServiceHandler aiEmbeddingServiceHandler;

    @Bean
    public JedisPooled getJedisPooled() {
        return new JedisPooled(getUri());
    }

    @SneakyThrows
    public void embedding(List<Document> documents, AiCollection aiCollection) {
        VectorStore vectorStore = getVectorStore(aiCollection);
        if (vectorStore instanceof InitializingBean store) {
            store.afterPropertiesSet();
        }
        vectorStore.add(documents);
    }

    @SneakyThrows
    public List<Document> similaritySearch(String query, AiCollection aiCollection) {
        VectorStore vectorStore = getVectorStore(aiCollection);
        if (vectorStore instanceof InitializingBean store) {
            store.afterPropertiesSet();
        }
        return vectorStore.similaritySearch(query);
    }

    public VectorStore getVectorStore(AiCollection aiCollection) {
        AiModel aiModel = aiCollection.embeddingModel();
        AiTag aiTag = aiModel.tagsView().stream()
                .filter(tag -> BaseAiModelTag.EMBEDDINGS.equals(tag.modelTag()))
                .findFirst()
                .orElseThrow(() -> new BusinessException("模型未配置Embeddings标签"));
        AiEmbeddingService aiEmbeddingService = aiEmbeddingServiceHandler.getByFactory(aiTag.factory())
                .orElseThrow(() -> new BusinessException("暂不支持该模型"));
        Map<String, Object> options = new HashMap<>(aiCollection.embeddingModel().options());
        EmbeddingModel embeddingModel = aiEmbeddingService.getEmbeddingModel(options);
        return new RedisVectorStore(getConfig(aiCollection.collectionName()), embeddingModel,
                getJedisPooled(), true);
    }

    public RedisVectorStore.RedisVectorStoreConfig getConfig(String collectionName) {
        return RedisVectorStore.RedisVectorStoreConfig.builder()
                .withIndexName(collectionName)
                .build();
    }

    public String getUri() {
        String username = StringUtils.hasText(redisConnectionDetails.getUsername()) ? redisConnectionDetails.getUsername() : "default";
        return "redis://" + username + ":" +
                redisConnectionDetails.getPassword() + "@" +
                redisConnectionDetails.getStandalone().getHost() + ":" +
                redisConnectionDetails.getStandalone().getPort();
    }

    @Override
    public void deleteCollection(String collectionName) {
        getJedisPooled().ftDropIndex(collectionName);
    }

    @Override
    public VectorType getVectorType() {
        return LpVectorType.REDIS;
    }
}
